import os
from typing import Annotated

from langchain_community.chat_models import ChatZhipuAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from typing_extensions import TypedDict

os.environ["ZHIPUAI_API_KEY"] = "97738d4998b8732d707daf91a2b1c56d.2y6VKEuOlidwHDpI"
os.environ["TAVILY_API_KEY"] = "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"

tool = TavilySearchResults(max_results=2)
tool_node = ToolNode(tools=[tool])
tools = [tool]

memory = MemorySaver()


class State(TypedDict):
    messages: Annotated[list, add_messages]


graph_builder = StateGraph(State)

llm = ChatZhipuAI(
    model="glm-4",
    temperature=0.5,
).bind_tools(tools)


def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}


graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", tool_node)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_conditional_edges("chatbot", tools_condition)

config = {"configurable": {"thread_id": "1"}}
graph = graph_builder.compile(checkpointer=memory)


def stream_graph_updates(user_input: str, thread_id: str):
    for event in graph.stream(
            {"messages": [("user", user_input)]},
            config
    ):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)


while True:
    user_input = input("User:")
    thread_id = input("ThreadId:")
    if user_input.lower() in ['quit', 'exit', 'q']:
        print("Goodbye!")
        break
    stream_graph_updates(user_input, thread_id)
    print(graph.get_state(config).values["messages"])
